Let AOTAutogradCache cache graphs with custom autograd.Functions - #2162

UMA's forward contains custom autograd.Functions. AOTAutogradCache refuses to
cache any graph containing one unless
torch._functorch.config.autograd_cache_allow_custom_autograd_functions is set,
so today the largest graph misses the cache and is retraced and recompiled from
scratch on every process start. Warm start ends up paying most of a cold
compile.

Turning it on halves warm compile time. On uma-s-1p2 with 1000 atoms, a second
process run against caches populated by the first:

  autograd_cache_allow_custom_autograd_functions=False   22.6s
  autograd_cache_allow_custom_autograd_functions=True    10.2s

which is a >2x reduction (3 runs each, 22.79/22.35/22.69 vs 10.86/9.92/9.76,
measured as the union of top-level compile events from tlparse). Cold start is
unaffected, 53.5s vs 54.0s, within run-to-run noise: the flag decides only
whether the second run can reuse the first run's work.

PyTorch intends to flip this flag to True by default, but that is waiting on
thorough internal testing across models. Nothing about UMA needs to wait for it:
the model's custom autograd.Functions are ordinary differentiable ops with no
state that would make a cached graph unsafe to reuse, so we can opt in now and
drop this line once the default changes.

This change was authored with the assistance of an AI coding agent.